import angr
import argparse

def get_address_space(binary_path):
    project = angr.Project(binary_path, auto_load_libs=False)
    main_obj = project.loader.main_object
    
    # 主二进制文件的最小和最大地址
    min_addr = main_obj.min_addr
    max_addr = main_obj.max_addr

    return min_addr, max_addr

def analyses_cfg(binary_path, min_addr, max_addr):
    # 设置超时
        # 初始化项目
    project = angr.Project(binary_path, auto_load_libs=False)

    # 定义初始变量
    function_addrs = set()  # 记录所有已分析的函数地址
    unique_addrs = set()  # 用于去重节点地址
    newly_discovered_addrs = set([project.entry])  # 初始化从入口点开始 project.entry 0x4020d0
    call_depth = 4  # 初始分析的深度

    cfg = None  # 存储最完整的 CFG
    temp_cfg = None

    # 持续分析直到没有新的函数发现
    while newly_discovered_addrs:
        # print(f"开始新一轮分析，分析深度: {call_depth}")
        current_discovered_addrs = set()  # 当前轮次新发现的函数地址

        # 对每个未分析的函数地址进行分析
        for addr in newly_discovered_addrs:
            if min_addr <= addr <= max_addr and addr not in unique_addrs:    
                # print(f"分析函数地址：{hex(addr)}")
                temp_cfg = project.analyses.CFGEmulated(
                    context_sensitivity_level=2,
                    keep_state=True,
                    call_depth=call_depth,
                    starts=[addr],  # 从特定函数地址开始分析
                    state_add_options=angr.sim_options.refs
                )
                unique_addrs.add(addr)

                # 如果是第一次执行分析，则保存此 CFG
                if cfg is None:
                    cfg = temp_cfg
                else:
                    # 将新 CFG 的函数、块、调用信息合并到现有的 CFG 中
                    # 手动合并函数，因为 FunctionManager 没有 update 方法
                    for func_addr, func in temp_cfg.kb.functions.items():
                        if func_addr not in cfg.kb.functions:
                            cfg.kb.functions[func_addr] = func

                # 记录新发现的函数地址
                for func in temp_cfg.kb.functions.values():
                    if func.addr not in function_addrs:
                        current_discovered_addrs.add(func.addr)

        # 更新已分析的函数地址
        function_addrs.update(newly_discovered_addrs)

        # 更新要分析的新函数地址
        newly_discovered_addrs = current_discovered_addrs - function_addrs
        # print(f"发现新的函数数量：{len(newly_discovered_addrs)}")

        # 如果没有新函数发现，可以选择加深分析深度
        if not newly_discovered_addrs:
            break

    return project, cfg

def generate_disassembly(binary_path, asm_path, block_path):

    min_addr, max_addr = get_address_space(binary_path)
    # 进行CFG分析
    # project, cfg = analyses_cfg(binary_path, min_addr, max_addr)
    # # 加载二进制文件
    project = angr.Project(binary_path, auto_load_libs=False)  # 加载二进制文件
    cfg = project.analyses.CFGFast()  # 获取 CFG
    
    # 打开输出文件
    with open(asm_path, 'w') as asm_file:
        # 遍历所有函数
        for func in cfg.kb.functions.values():
            # 判断函数是否在指定的地址范围内
            if func.addr < min_addr or func.addr > max_addr:
                continue  # 跳过不在地址范围内的函数

            # 写入函数名
            asm_file.write(f"function name:{func.name}\n")
            
            # 获取函数的地址和汇编代码
            for block in func.blocks:
                # 每个块的起始地址
                start_addr = block.addr
                # 获取该地址的汇编代码
                disassembly = block.capstone.insns
                for insn in disassembly:
                    # 写入地址和汇编代码
                    asm_file.write(f"0x{start_addr:X}: {insn.mnemonic} {insn.op_str}\n")
                    start_addr += len(insn.bytes)  # 更新地址为下一条指令的地址
            asm_file.write("\n")  # 每个函数之间空一行

    print(f"反汇编代码已保存到 {asm_path}")
    
    with open(block_path,"w") as block_file:
        block_file.write("\n")
        for func in cfg.kb.functions.values():
            if func.addr > max_addr or func.addr < min_addr :
                continue
            block_file.write(f"// Function: {func.name}\n")
            for block in func.blocks:
                end_addr = block.addr + block.size-1
                block_file.write(f"Block Start: 0x{block.addr:x}, Block End: 0x{end_addr:x}\n")
            block_file.write("\n")  # 添加空行以分隔不同基本块

# # 示例用法：分析指定的二进制文件并保存为反汇编文件
# binary_path = r"D:\Users\Administrator\Desktop\study\project\dataset\note_test\binary\bad\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82.out"  # 请替换为你要分析的二进制文件路径
# output_path = r"D:\Users\Administrator\Desktop\study\project\dataset\note_test\json\bad1\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82.asm"  # 输出文件路径，反汇编代码将保存在该文件中

# generate_disassembly(binary_path, output_path)

# # 示例用法：分析指定的二进制文件并保存为反汇编文件
# binary_path = r"D:\Users\Administrator\Desktop\study\项目\IDA\source_data\latex2rtf\latex2rtf_32"  # 请替换为你要分析的二进制文件路径
# output_path = r"D:\Users\Administrator\Desktop\study\项目\IDA\out_res\latex2rtf\latex2rtf_32.asm"  # 输出文件路径，反汇编代码将保存在该文件中

# generate_disassembly(binary_path, output_path)

def main():
    parser = argparse.ArgumentParser(description='Angr Disassembly Script')
    parser.add_argument('--binary_path', required=True, help='Path to the binary file')
    parser.add_argument('--asm_path', required=True, help='Path to save the disassembly output')
    parser.add_argument('--block_path', required=True, help='Path to save the block output')

    args = parser.parse_args()

    generate_disassembly(args.binary_path, args.asm_path, args.block_path)

if __name__ == "__main__":
    main()